import requests
import uuid
import yaml
import os, sys


class WebSearcher():
    def __init__(self, **kwargs):
        self.CURRENT_FOLDER = os.path.dirname(os.path.abspath(__file__))
        with open(os.path.join(self.CURRENT_FOLDER, 'config.yaml'), 'r', encoding='utf-8') as yaml_file:
            self.config = yaml.load(yaml_file, Loader=yaml.FullLoader)
        self.kwargs = kwargs
        self.api_key = os.environ.get("ZHIPUAI_API_KEY", self.config['zhipuai_api_key'])
        self.tool = "web-search-pro"
        self.url = "https://open.bigmodel.cn/api/paas/v4/tools"

    # 必须使用中文输入，才能得到有效搜索结果
    def search(self, user_query="what is loss function"):
        msg = [
            {
                "role": "user",
                "content": user_query
            }
        ]
        request_id = str(uuid.uuid4())
        data = {
            "request_id": request_id,
            "tool": self.tool,
            "stream": False,
            "messages": msg
        }

        resp = requests.post(
            self.url,
            json=data,
            headers={'Authorization': self.api_key},
            timeout=1000
        )
        print(resp.content.decode())



if __name__ == '__main__':
    web_searcher = WebSearcher()
    
    web_searcher.search("什么是损失函数？")